{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "color-singer",
   "metadata": {},
   "source": [
    "# Talking Head Anime from a Single Image (Manual Poser Tool)\n",
    "\n",
    "**Instruction**\n",
    "\n",
    "1. From the main menu, click \"Runtime > Change runtime type.\" \n",
    "2. Change \"Hardware accelerator\" to \"GPU,\" and click \"Save.\"\n",
    "3. Run the four cells below, one by one, in order by clicking the \"Play\" button to the left of it. Wait for each cell to finish before going to the next one.\n",
    "4. Scroll down to the end of the last cell, and play with the GUI.\n",
    "\n",
    "**Constraints on Images**\n",
    "\n",
    "1. Must be an image of a single humanoid anime character.\n",
    "2. Must be of size 256x256.\n",
    "3. The head must be roughly contained in the middle 128x128 middle box.\n",
    "4. Must have PNG format.\n",
    "5. Must have an alpha channel.\n",
    "6. Background pixels must have RGBA=(0,0,0,0). See [this link](https://github.com/pkhungurn/talking-head-anime-demo/issues/1) if you do not get clean results.\n",
    "\n",
    "**Links**\n",
    "\n",
    "* Github repository: http://github.com/pkhungurn/talking-head-anime-demo\n",
    "* Project writeup: http://pkhungurn.github.io/talking-head-anime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bottom-tuner",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clone the repository\n",
    "%cd /content\n",
    "!git clone https://github.com/pkhungurn/talking-head-anime-demo.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "significant-burlington",
   "metadata": {},
   "outputs": [],
   "source": [
    "# CD into the repository directory.\n",
    "%cd /content/talking-head-anime-demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "federal-youth",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download model files\n",
    "!wget -O data/combiner.pt https://www.dropbox.com/s/p220v9rmbjmqien/combiner.pt?dl=0\n",
    "!wget -O data/face_morpher.pt https://www.dropbox.com/s/oukbnofkffc2bis/face_morpher.pt?dl=0\n",
    "!wget -O data/two_algo_face_rotator.pt https://www.dropbox.com/s/o78wzc5cpxnxggr/two_algo_face_rotator.pt?dl=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "written-shift",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the GUI\n",
    "\n",
    "import torch\n",
    "\n",
    "DEVICE_NAME = 'cuda'\n",
    "device = torch.device(DEVICE_NAME)\n",
    "\n",
    "import PIL.Image\n",
    "import io\n",
    "from io import BytesIO\n",
    "import IPython.display\n",
    "import numpy\n",
    "import ipywidgets\n",
    "from poser.morph_rotate_combine_poser import MorphRotateCombinePoser256Param6\n",
    "from poser.poser import Poser\n",
    "from tha.combiner import CombinerSpec\n",
    "from tha.face_morpher import FaceMorpherSpec\n",
    "from tha.two_algo_face_rotator import TwoAlgoFaceRotatorSpec\n",
    "from util import extract_pytorch_image_from_filelike, rgba_to_numpy_image\n",
    "\n",
    "last_torch_input_image = None\n",
    "torch_input_image = None\n",
    "\n",
    "def show_pytorch_image(pytorch_image, output_widget=None):\n",
    "    output_image = pytorch_image.detach().cpu()\n",
    "    numpy_image = rgba_to_numpy_image(output_image)    \n",
    "    pil_image = PIL.Image.fromarray(numpy.uint8(numpy.rint(numpy_image * 255.0)), mode='RGBA')        \n",
    "    IPython.display.display(pil_image)\n",
    "\n",
    "input_image_widget = ipywidgets.Output(\n",
    "    layout={\n",
    "        'border': '1px solid black',\n",
    "        'width': '256px',\n",
    "        'height': '256px'\n",
    "    })\n",
    "\n",
    "upload_input_image_button = ipywidgets.FileUpload(\n",
    "    accept='.png',\n",
    "    multiple=False,\n",
    "    layout={\n",
    "        'width': '256px'\n",
    "    }\n",
    ")\n",
    "\n",
    "output_image_widget = ipywidgets.Output(\n",
    "    layout={\n",
    "        'border': '1px solid black',\n",
    "        'width': '256px',\n",
    "        'height': '256px'\n",
    "    }\n",
    ")\n",
    "\n",
    "eye_left_slider = ipywidgets.FloatSlider(\n",
    "    value=0.0,\n",
    "    min=0.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    description=\"Left Eye:\",\n",
    "    readout=True,\n",
    "    readout_format=\".2f\"\n",
    ")\n",
    "eye_right_slider = ipywidgets.FloatSlider(\n",
    "    value=0.0,\n",
    "    min=0.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    description=\"Right Eye:\",\n",
    "    readout=True,\n",
    "    readout_format=\".2f\"\n",
    ")\n",
    "mouth_slider = ipywidgets.FloatSlider(\n",
    "    value=0.0,\n",
    "    min=0.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    description=\"Mouth:\",\n",
    "    readout=True,\n",
    "    readout_format=\".2f\"\n",
    ")\n",
    "\n",
    "head_x_slider = ipywidgets.FloatSlider(\n",
    "    value=0.0,\n",
    "    min=-1.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    description=\"X-axis:\",\n",
    "    readout=True,\n",
    "    readout_format=\".2f\"\n",
    ")\n",
    "head_y_slider = ipywidgets.FloatSlider(\n",
    "    value=0.0,\n",
    "    min=-1.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    description=\"Y-axis:\",\n",
    "    readout=True,\n",
    "    readout_format=\".2f\",    \n",
    ")\n",
    "neck_z_slider = ipywidgets.FloatSlider(\n",
    "    value=0.0,\n",
    "    min=-1.0,\n",
    "    max=1.0,\n",
    "    step=0.01,\n",
    "    description=\"Z-axis:\",\n",
    "    readout=True,\n",
    "    readout_format=\".2f\",    \n",
    ")\n",
    "\n",
    "\n",
    "control_panel = ipywidgets.VBox([    \n",
    "    ipywidgets.HTML(value=\"<center><b>Head Rotation</b></center>\"),\n",
    "    head_x_slider,\n",
    "    head_y_slider,\n",
    "    neck_z_slider,\n",
    "    ipywidgets.HTML(value=\"<hr>\"),\n",
    "    ipywidgets.HTML(value=\"<center><b>Facial Features</b></center>\"),\n",
    "    eye_left_slider,\n",
    "    eye_right_slider,\n",
    "    mouth_slider,\n",
    "])\n",
    "\n",
    "controls = ipywidgets.HBox([\n",
    "    ipywidgets.VBox([\n",
    "        input_image_widget, \n",
    "        upload_input_image_button\n",
    "    ]),\n",
    "    control_panel,\n",
    "    ipywidgets.HTML(value=\"&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;\"),\n",
    "    output_image_widget,\n",
    "])\n",
    "\n",
    "poser = MorphRotateCombinePoser256Param6(\n",
    "    morph_module_spec=FaceMorpherSpec(),\n",
    "    morph_module_file_name=\"data/face_morpher.pt\",\n",
    "    rotate_module_spec=TwoAlgoFaceRotatorSpec(),\n",
    "    rotate_module_file_name=\"data/two_algo_face_rotator.pt\",\n",
    "    combine_module_spec=CombinerSpec(),\n",
    "    combine_module_file_name=\"data/combiner.pt\",\n",
    "    device=device)\n",
    "pose_size = 6\n",
    "last_pose = torch.zeros(1, pose_size).to(device)\n",
    "\n",
    "def get_pose():\n",
    "    pose = torch.zeros(1, pose_size)\n",
    "    \n",
    "    pose[0, 0] = head_x_slider.value\n",
    "    pose[0, 1] = head_y_slider.value\n",
    "    pose[0, 2] = neck_z_slider.value\n",
    "    pose[0, 3] = eye_left_slider.value\n",
    "    pose[0, 4] = eye_right_slider.value\n",
    "    pose[0, 5] = mouth_slider.value\n",
    "        \n",
    "    return pose.to(device)\n",
    "\n",
    "display(controls)\n",
    "\n",
    "def update(change):\n",
    "    global last_pose\n",
    "    global last_torch_input_image\n",
    "        \n",
    "    if torch_input_image is None:\n",
    "        return\n",
    "        \n",
    "    needs_update = False\n",
    "    if last_torch_input_image is None:\n",
    "        needs_update = True        \n",
    "    else:\n",
    "        if (torch_input_image - last_torch_input_image).abs().max().item() > 0:\n",
    "            needs_update = True         \n",
    "            \n",
    "    pose = get_pose()\n",
    "    if (pose - last_pose).abs().max().item() > 0:\n",
    "        needs_update = True\n",
    "    \n",
    "    if not needs_update:\n",
    "        return\n",
    "   \n",
    "    output_image = poser.pose(torch_input_image, pose)[0]\n",
    "    with output_image_widget:\n",
    "        output_image_widget.clear_output(wait=True)\n",
    "        show_pytorch_image(output_image, output_image_widget)  \n",
    "        \n",
    "    last_torch_input_image = torch_input_image\n",
    "    last_pose = pose\n",
    "        \n",
    "def upload_image(change):\n",
    "    global torch_input_image\n",
    "    for name, file_info in upload_input_image_button.value.items():\n",
    "        torch_input_image = extract_pytorch_image_from_filelike(io.BytesIO(file_info['content'])).to(device)\n",
    "        torch_input_image = torch_input_image.unsqueeze(0)\n",
    "    if torch_input_image is not None:\n",
    "        n,c,h,w = torch_input_image.shape\n",
    "        if h != 256 or w != 256:\n",
    "            with input_image_widget:\n",
    "                input_image_widget.clear_output(wait=True)\n",
    "                display(ipywidgets.HTML(\"Image must be 256x256 in size!!!\"))\n",
    "            torch_input_image = None\n",
    "        if c != 4:\n",
    "            with input_image_widget:\n",
    "                input_image_widget.clear_output(wait=True)\n",
    "                display(ipywidgets.HTML(\"Image must have an alpha channel!!!\"))                \n",
    "            torch_input_image = None\n",
    "        if torch_input_image is not None:\n",
    "            with input_image_widget:\n",
    "                input_image_widget.clear_output(wait=True)\n",
    "                show_pytorch_image(torch_input_image[0], input_image_widget)\n",
    "        update(None)\n",
    "        \n",
    "upload_input_image_button.observe(upload_image, names='value')\n",
    "eye_left_slider.observe(update, 'value')\n",
    "eye_right_slider.observe(update, 'value')\n",
    "mouth_slider.observe(update, 'value')\n",
    "head_x_slider.observe(update, 'value')\n",
    "head_y_slider.observe(update, 'value')\n",
    "neck_z_slider.observe(update, 'value')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
